package com.gitee.feizns.quickstart.web.advice.crypto;

import com.gitee.feizns.dynamic.Maps;
import com.gitee.feizns.dynamic.decrypt.Aes;
import com.gitee.feizns.dynamic.decrypt.Rsa;
import com.gitee.feizns.dynamic.io.Ios;
import com.gitee.feizns.quickstart.jackson.JsonUtils;
import com.gitee.feizns.quickstart.web.WebUtils;
import com.gitee.feizns.quickstart.web.decrypt.DecryptionAlgorithm;
import com.gitee.feizns.quickstart.web.ex.UnauthorizedException;
import lombok.Getter;
import lombok.SneakyThrows;
import org.springframework.core.MethodParameter;
import org.springframework.core.Ordered;
import org.springframework.core.annotation.Order;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpInputMessage;
import org.springframework.http.MediaType;
import org.springframework.http.converter.HttpMessageConverter;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.lang.NonNull;
import org.springframework.util.AntPathMatcher;
import org.springframework.web.bind.annotation.RestControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.RequestBodyAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.lang.reflect.Type;
import java.nio.charset.StandardCharsets;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.Objects;

/**
 * 加解密前后置处理器
 * 启动加解密：
 * 请求：携带请求头：Qk-Crypto-Pub-Key=RSA公钥（响应体密钥加密公钥）、Qk-Crypto-Key=AES密钥用后端公钥加密后的密钥
 * 响应：携带响应头：Qk-Crypto-Key（用于响应体解密的AES密钥）
 * 前后端各准备一对RSA和AES密钥
 * @author feizns
 * @since 2025/9/25
 */
@Order(Ordered.HIGHEST_PRECEDENCE)
@RestControllerAdvice
public class CryptoRequestResponseBodyAdvice implements RequestBodyAdvice, ResponseBodyAdvice<Object> {

    /**
     * 匹配器
     */
    private final AntPathMatcher matcher = new AntPathMatcher();

    /**
     * 媒体类型
     */
    @Getter
    private final List<MediaType> mediaTypes;

    /**
     * 路径
     */
    @Getter
    private final List<String> patterns;

    /**
     * 加密请求响应正文建议
     */
    public CryptoRequestResponseBodyAdvice() {
        this(Collections.emptyList(), Collections.emptyList());
    }

    /**
     * 加密请求响应正文建议
     */
    public CryptoRequestResponseBodyAdvice(List<MediaType> mediaTypes) {
        this(mediaTypes, Collections.emptyList());
    }

    /**
     * 加密请求响应正文建议
     */
    public CryptoRequestResponseBodyAdvice(List<MediaType> mediaTypes, List<String> patterns) {
        this.mediaTypes = mediaTypes;
        this.patterns = patterns;
    }

    /**
     * 获取加密发布密钥
     * @return {@link String }
     */
    public Map.Entry<String, String> getCryptoPubKey() {
        String qkCryptoPubKey = WebUtils.getReq().getHeader("Qk-Crypto-Pub-Key");
        if ( qkCryptoPubKey == null )
            throw new UnauthorizedException("Qk-Crypto-Pub-Key不能为空.");
        return Maps.entry("Qk-Crypto-Pub-Key", qkCryptoPubKey);
    }

    /**
     * 获取加密密钥
     * @return {@link String }
     */
    public Map.Entry<String, String> getCryptoKey() {
        String qkCryptoPubKey = WebUtils.getReq().getHeader("Qk-Crypto-Key");
        if ( qkCryptoPubKey == null )
            throw new UnauthorizedException("Qk-Crypto-Key不能为空.");
        return Maps.entry("Qk-Crypto-Key", qkCryptoPubKey);
    }

    /**
     * 解密
     * @param body 身体
     * @return {@link String }
     */
    public String decrypt(InputStream body) {
        return decrypt(Ios.asString(body));
    }

    /**
     * 解密
     * @param body 身体
     * @return {@link String }
     */
    public String decrypt(String body) {
        Map.Entry<String, String> cryptoKey = getCryptoKey();
        String aesKey = DecryptionAlgorithm.instance().decrypt(cryptoKey.getValue());
        return Aes.decrypt(body, aesKey);
    }

    /**
     * 加密
     * @param plain 平原
     * @return {@link String }
     */
    public String encrypt(String plain) {
        Map.Entry<String, String> cryptoPubKey = getCryptoPubKey();
        String aesKey = Aes.generateBase64SecretKey();
        String cipher = Aes.encrypt(plain, aesKey);
        String cryptoKey = Rsa.instance().encrypt(aesKey, cryptoPubKey.getValue());
        WebUtils.getResp().addHeader(cryptoPubKey.getKey(), cryptoKey);
        return cipher;
    }

    @Override
    public boolean supports(@NonNull MethodParameter methodParameter, @NonNull Type targetType, @NonNull Class<? extends HttpMessageConverter<?>> converterType) {
        return isCrypto(methodParameter);
    }

    @Override
    public @NonNull HttpInputMessage beforeBodyRead(@NonNull HttpInputMessage inputMessage, @NonNull MethodParameter parameter, @NonNull Type targetType, @NonNull Class<? extends HttpMessageConverter<?>> converterType) throws IOException {
        String plain = decrypt(inputMessage.getBody());
        return new HttpInputMessage() {
            @Override
            public @NonNull InputStream getBody() {
                return new ByteArrayInputStream(plain.getBytes(StandardCharsets.UTF_8));
            }

            @Override
            public @NonNull HttpHeaders getHeaders() {
                return inputMessage.getHeaders();
            }
        };
    }

    @Override
    public @NonNull Object afterBodyRead(@NonNull Object body, @NonNull HttpInputMessage inputMessage, @NonNull MethodParameter parameter, @NonNull Type targetType, @NonNull Class<? extends HttpMessageConverter<?>> converterType) {
        return body;
    }

    @Override
    public Object handleEmptyBody(Object body, @NonNull HttpInputMessage inputMessage, @NonNull MethodParameter parameter, @NonNull Type targetType, @NonNull Class<? extends HttpMessageConverter<?>> converterType) {
        return body;
    }

    @Override
    public boolean supports(@NonNull MethodParameter returnType, @NonNull Class<? extends HttpMessageConverter<?>> converterType) {
        return isCrypto(returnType);
    }

    @SneakyThrows
    @Override
    public Object beforeBodyWrite(Object body, @NonNull MethodParameter returnType, @NonNull MediaType selectedContentType, @NonNull Class<? extends HttpMessageConverter<?>> selectedConverterType, @NonNull ServerHttpRequest request, @NonNull ServerHttpResponse response) {
        return encrypt(JsonUtils.toJson(body));
    }

    /**
     * 是否是加密
     * @param methodParameter 方法参数
     * @return boolean
     */
    public boolean isCrypto(MethodParameter methodParameter) {
        if ( mediaTypes.stream().anyMatch(type -> WebUtils.getMediaType().isCompatibleWith(type)) ) {
            return methodParameter.hasMethodAnnotation(CryptoBody.class)
                    || Objects.requireNonNull(methodParameter.getMethod()).getDeclaringClass().isAnnotationPresent(CryptoBody.class)
                    || (patterns.stream().anyMatch(pattern -> matcher.match(pattern, WebUtils.getRequestURI())));
        }
        return false;
    }

}
